# -*- coding: utf-8 -*-

import numpy
import cv2
from splitIMG import *

file_dir = '../../images'
class_type = 'plane'
cut_path = '../../process/'
cut_label_path = 'label/'
index = 0
if not os.path.exists(cut_path): 
    os.mkdir(cut_path)
    os.mkdir(cut_path + cut_label_path)

root, patches = getPatches(file_dir)
for patch in patches:
    labels = loadLabels(root['label'] + patch + '.txt', class_type)
    if not labels: continue
    image = cv2.imread(root['image'] + patch + '.png')
    crops = cutImage(image)
    for crop in crops:
        new_labels = []

        for label in labels:
            in_crop = (min(label[::2]) > crop[0]) \
                    and (min(label[1::2]) > crop[1]) \
                    and (max(label[::2]) < crop[2]) \
                    and (max(label[1::2]) < crop[3]) 
            if not in_crop: continue
            
            new_labels.append([
                label[0] - crop[0],
                label[1] - crop[1],
                label[2] - crop[0],
                label[3] - crop[1],
                label[4] - crop[0],
                label[5] - crop[1],
                label[6] - crop[0],
                label[7] - crop[1]
            ])
        
        if not new_labels: continue
        print(crop)
        crop_img = image[crop[1]:crop[3], crop[0]:crop[2]]
        cv2.imwrite(cut_path + str(index) + '.png', crop_img)
        new_label_file = open(cut_path + cut_label_path + str(index) + '.txt', 'a+')
        for line in new_labels:
            for num in line:
                new_label_file.write(str(num) + ' ')
            new_label_file.write('plane \n')
        new_label_file.close()
        index += 1